import yaml
import psycopg2
from typing import Dict, Any, Optional
from psycopg2 import OperationalError, InterfaceError, connect as pg_connect
import re
from sqlalchemy.engine import create_engine, Engine

class DatabaseConfig:
    def __init__(self, config_file: str):
        """
        初始化PostgreSQL DatabaseConfig

        参数:
            config_file: YAML配置文件的路径
        """
        self.config_file = config_file
        self.connection = None
        self.config = self._parse_config()
        self.dsn = self.get_connstr()
    def get_connstr(self)->str:
        return "postgresql+psycopg2://{}:{}@{}:{}/{}".format(
            self.config["user"],
            self.config["password"],
            self.config["host"],
            self.config["port"],
            self.config["dbname"])
    def _validate_config(self, config: Dict[str, Any]) -> Optional[str]:
        """
        验证PostgreSQL连接配置参数的有效性
        参数:
            config: 连接配置字典
        返回:
            如果验证通过返回None，否则返回错误信息字符串
        """
        # 必填字段检查
        required_fields = ['dbname', 'user']
        for field in required_fields:
            if not config.get(field):
                return f"缺少必填字段: {field}"
        
        # 主机名或IP检查
        host = config.get('host', 'localhost')
        if not (re.match(r'^[a-zA-Z0-9\-\.]+$', host) or host == 'localhost'):
            return "主机名格式无效"
        
        # 端口检查
        port = config.get('port', 5432)
        if not (1 <= int(port) <= 65535):
            return "端口号必须在1-65535范围内"
        
        # SSL模式检查
        valid_ssl_modes = ['disable', 'allow', 'prefer', 'require', 'verify-ca', 'verify-full']
        sslmode = config.get('sslmode', 'prefer')
        if sslmode not in valid_ssl_modes:
            return f"无效的sslmode，必须是: {', '.join(valid_ssl_modes)}"
        return None
        
    def _parse_config(self) -> Dict[str, Any]:
        """
        解析包含PostgreSQL连接配置的YAML文件并进行基本校验
        
        参数:
            yaml_file: YAML配置文件的路径
            
        返回:
            包含PostgreSQL连接参数的字典
            
        抛出:
            ValueError: 当配置无效时
            
        示例YAML格式:
            postgresql:
              host: "localhost"
              port: 5432
              database: "mydb"
              user: "myuser"
              password: "mypassword"
              options:
                - "-c search_path=public"
              sslmode: "prefer"
        """
        try:
            with open(self.config_file, 'r') as f:
                config = yaml.safe_load(f)
                
            # 获取postgresql部分的配置
            pg_config = config.get('postgres', {})
            # print(pg_config)
            # 构造连接参数字典
            connection_params = {
                'host': pg_config.get('host', 'localhost'),
                'port': int(pg_config.get('port', 5432)),
                'dbname': pg_config.get('database', ''),
                'user': pg_config.get('user', ''),
                'password': pg_config.get('password', ''),
                'sslmode': pg_config.get('sslmode', 'prefer')
            }
            
            # 处理额外的连接选项
            if 'options' in pg_config:
                if isinstance(pg_config['options'], list):
                    connection_params['options'] = ' '.join(pg_config['options'])
                else:
                    raise ValueError("options必须是列表格式")
                
            # 移除空值的参数（除了password，因为可以是空密码）
            connection_params = {k: v for k, v in connection_params.items() 
                               if v or k == 'password'}
            
            # 验证配置
            validation_error = self._validate_config(connection_params)
            if validation_error:
                raise ValueError(f"配置验证失败: {validation_error}")
            
            return connection_params
            
        except FileNotFoundError:
            raise ValueError(f"配置文件 {self.config_file} 未找到")
        except yaml.YAMLError as e:
            raise ValueError(f"YAML解析错误: {e}")
        except ValueError as e:
            raise e
        except Exception as e:
            raise ValueError(f"解析配置时出错: {e}")
  

class DatabaseManager(DatabaseConfig):
    def __init__(self, config_file: str):
        super().__init__(config_file)
        self.provider = "postgres"
        self.connection= self._connect()
    def set_provider(self, provider: str):
        self.provider = provider 
    def create_engine(self) -> Engine:
        if self.provider == "opengauss":
            print("openGauss")
            from sqlalchemy.dialects.postgresql.base import PGDialect
            PGDialect._get_server_version_info = lambda *args: (9, 2)
            return create_engine(self.dsn, pool_pre_ping=True)
        return create_engine(self.dsn, pool_pre_ping=True)
    def connect(self) -> psycopg2.extensions.connection:
        return self.connection
    def _connect(self) -> psycopg2.extensions.connection:
        """
        连接到PostgreSQL数据库
        
        返回:
            psycopg2连接对象
            
        抛出:
            ConnectionError: 当连接失败时，包含详细错误信息
        """
        try:
            self._connection = pg_connect(**self.config)
            return self._connection
        except OperationalError as e:
            raise ConnectionError(f"连接操作失败: {e}")
        except InterfaceError as e:
            raise ConnectionError(f"连接接口错误: {e}")
        except Exception as e:
            raise ConnectionError(f"连接时发生意外错误: {e}")
    def close(self):
        """关闭数据库连接"""
        if self._connection and not self._connection.closed:
            self._connection.close()
        self._connection = None